#!/usr/bin/env bash

set -x

setup_env() {
    set -xeuo pipefail

    . conda-dnld-utils
    setup_conda_channel;
    export DEBIAN_FRONTEND=non-interactive

    # Run package updates and install packages
    apt-get update
    apt-get install -y numactl make

    legate-conda-retry search --override-channels -c "${ARTIFACTS_DIR}/conda-build/cupynumeric" --info cupynumeric

    # This requires strict channel priority to work (prioritize local channel)
    legate-conda-retry create -y -n legate -c "${ARTIFACTS_DIR}/conda-build/cupynumeric" -c https://conda.anaconda.org/${CONDA_CHANNEL}/label/${CONDA_LABEL} -c conda-forge legate cupynumeric
}

setup_test_env() {
    legate-conda-retry install -y "clang-tools>=8" "clang>=8" colorama coverage mock pre-commit pytest-cov pytest-mock "pytest" types-docutils pynvml psutil

    if [[ "${CUPY_REQUIRED_FOR_TESTS:-0}" == "1" ]]; then
        legate-conda-retry install -y cupy
    fi

    pip install tifffile
}

setup_docs_env() {
    legate-conda-retry install -y pandoc doxygen
    pip install ipython jinja2 "markdown<3.4.0" myst-parser  nbsphinx  sphinx-copybutton  "sphinx>=8"  nvidia-sphinx-theme cffi
}

setup_mypy_env() {
    legate-conda-retry install -y "mypy>=0.961" jinja2 nbsphinx sphinx-copybutton  "sphinx>=4.4.0" types-docutils scipy-stubs
}

setup_unit_env() {
    legate-conda-retry install -y pytest pytest-mock mock cffi
}

run_legate_issue() {
    if command -v "legate-issue" &> /dev/null; then
      legate-issue
    else
      echo "WARNING: legate-issue not found."
    fi
}

test_cupynumeric() {
    set -xeo pipefail

    . conda-utils;
    export CONDA_ENV=legate

    setup_env;

    activate_conda_env;
    conda_info;

    cd "${REPO_DIR}";

    export WORKERS=""
    if command -v nvidia-smi &> /dev/null; then
        gpu_count=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
        if [ "$gpu_count" -ge 2 ]; then
            export WORKERS="-j 1"
        fi
    fi

    case "$1" in
        "test")
            echo "Executing tests..."
            shift;
            CUPY_REQUIRED_FOR_TESTS=0
            if [[ " $* " == *" --use cuda "* ]]; then
                CUPY_REQUIRED_FOR_TESTS=1
            fi
            setup_test_env;
            run_legate_issue;
            ./test.py ${WORKERS} -vv --timeout 600 "$@"
            ;;
        "mypy")
            echo "Installing and executing mypy..."
            shift;
            setup_mypy_env;
            run_legate_issue;
            mypy cupynumeric
            ;;
        "docs")
            echo "Building docs..."
            shift;
            setup_docs_env;
            run_legate_issue;
            cd docs/cupynumeric
            make clean html
            # ls -lah .
            echo Copying artifacts
            cd build/html
            cp -r . "${OUTPUT_ARTIFACTS_DIR}"
            ;;
        "unit")
            echo "Running Unit tests..."
            shift;
            setup_unit_env;
            run_legate_issue;
            LEGATE_AUTO_CONFIG=0 pytest tests/unit
            ;;
        "cpp")
            echo "Running CPP tests..."
            shift;
            run_legate_issue;
            export LD_LIBRARY_PATH=${CONDA_PREFIX}/lib/legate/deps:${LD_LIBRARY_PATH:-}
            REALM_BACKTRACE=1 LEGATE_TEST=1 ${CONDA_PREFIX}/bin/cpp_tests
            ;;
        *)
            echo "Invalid command: $1"
            return 1
            ;;
    esac
}

(test_cupynumeric "$@");
